"""
BTC/USDT 交易策略 - 分批入场版本
策略描述：
1. 短均线：3天收盘价SMA
2. 长均线：40天收盘价SMA
3. 均线多头排列：计算3-101区间的SMA，比较相邻短区间和长区间SMA
4. 分批入场，一次性出场
"""

from dataclasses import dataclass
import datetime
import pandas as pd
import numpy as np
import warnings
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")

try:
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
except:
    pass


@dataclass
class Performance:
    """用于记录策略表现"""
    annualized_return: float = 0.0
    cumulative_return: float = 0.0
    excess_return: float = 0.0
    excess_cumulative: float = 0.0
    max_drawdown: float = 0.0
    total_trades: float = 0.0
    yearly_trades: float = 0.0
    sharpe: float = 0.0
    calmar: float = 0.0
    winning_rate: float = 0.0
    annualized_volatility: float = 0.0
    annualized_sharpe: float = 0.0
    yearly_count: float = 0.0
    
    def dict(self):
        """转换为字典"""
        return {
            'annualized_return': self.annualized_return,
            'cumulative_return': self.cumulative_return,
            'excess_return': self.excess_return,
            'excess_cumulative': self.excess_cumulative,
            'max_drawdown': self.max_drawdown,
            'total_trades': self.total_trades,
            'yearly_trades': self.yearly_trades,
            'sharpe': self.sharpe,
            'calmar': self.calmar,
            'winning_rate': self.winning_rate,
            'annualized_volatility': self.annualized_volatility,
            'annualized_sharpe': self.annualized_sharpe,
            'yearly_count': self.yearly_count
        }


class BTCTradingStrategy:
    """BTC/USDT 分批入场策略"""
    
    def __init__(self):
        self.frequency = "1d"
        self.data_begin_date = datetime.date(2015, 1, 1)
        self.backtest_begin_date = datetime.date(2020, 1, 1)
        self.backtest_end_date = datetime.date(2025, 9, 30)
        self.short_ma = 3
        self.long_ma = 40
        self.ma_range_start = 3
        self.ma_range_end = 101
        self.bullish_threshold = 0.6
        self.trading_days_one_year = 365
        self.risk_free_rate = 2.0
        self.btc_df = None
        self.equity_curve_df = None
        self.trade_log_df = None
        self.performance_overall = None
        self.performance_by_year = {}
        
    def load_data(self):
        """加载数据"""
        self.btc_df = pd.read_csv("BTC_kline_1d_merge.csv")
        self.btc_df["datetime"] = pd.to_datetime(self.btc_df["datetime"])
        self.btc_df["date"] = self.btc_df["datetime"].dt.date
        self.btc_df = self.btc_df[self.btc_df['date'] >= self.data_begin_date]
        self.btc_df = self.btc_df[self.btc_df['date'] <= self.backtest_end_date]
        print(f"数据加载完成，共{len(self.btc_df)}条记录")
        
    def calculate_ma_alignment(self):
        """计算群体多头排列比率"""
        print("计算群体多头排列指标...")
        
        def inner_func(row):
            if any(pd.isna(row)):
                return np.nan
            total_pairs = 0
            bullish_pairs = 0
            for i in range(1, len(row)):
                if row[i] < row[i-1]:  # 短均线>长均线
                    bullish_pairs += 1
                total_pairs += 1
            return bullish_pairs / total_pairs if total_pairs > 0 else 0
        
        ma_values = []
        for i in range(self.ma_range_start, self.ma_range_end):
            ma_values.append(self.btc_df["close"].rolling(window=i).mean())
        
        ma_all = pd.concat(ma_values, axis=1)
        self.btc_df["bullish_ratio"] = ma_all.apply(lambda row: inner_func(row), axis=1)
        print("群体多头排列指标计算完成")
        
    def execute_strategy(self):
        """执行分批入场策略"""
        print("开始执行策略...")
        
        self.btc_df["ma_s"] = self.btc_df["close"].rolling(window=self.short_ma).mean()
        self.btc_df["ma_l"] = self.btc_df["close"].rolling(window=self.long_ma).mean()
        
        truncate_btc_df = self.btc_df[self.btc_df['date'] >= self.backtest_begin_date].copy()
        
        # 使用更简单的逻辑
        position_ratio = 0.0  # 仓位比例：0, 0.5, 1.0
        pre_price = 0
        entry_price = 0
        equity = [10000]
        trade_log = []
        
        for idx, row in truncate_btc_df.iterrows():
            current_price = row["close"]
            bullish_ratio = row.get("bullish_ratio", 0)
            ma_s = row.get("ma_s", 0)
            ma_l = row.get("ma_l", 0)
            
            if pd.isna(ma_s) or pd.isna(ma_l) or pd.isna(bullish_ratio):
                equity.append(equity[-1])
                if pre_price == 0:
                    pre_price = current_price
                continue
            
            # 交易逻辑
            if position_ratio == 0 and ma_s > ma_l:
                # 空仓时，短均线突破长均线
                if bullish_ratio >= self.bullish_threshold:
                    # 直接买入100%
                    entry_price = current_price
                    position_ratio = 1.0
                    trade_log.append({
                        'datetime': row["datetime"],
                        'date': row["date"],
                        'action': 'BUY_100%',
                        'price': current_price
                    })
                else:
                    # 买入50%
                    entry_price = current_price
                    position_ratio = 0.5
                    trade_log.append({
                        'datetime': row["datetime"],
                        'date': row["date"],
                        'action': 'BUY_50%',
                        'price': current_price
                    })
                equity.append(equity[-1])
                    
            elif position_ratio == 0.5 and bullish_ratio >= self.bullish_threshold:
                # 从50%加仓到100%
                position_ratio = 1.0
                entry_price = current_price
                equity.append(equity[-1])
                trade_log.append({
                    'datetime': row["datetime"],
                    'date': row["date"],
                    'action': 'ADD_50%',
                    'price': current_price
                })
                
            elif position_ratio > 0 and ma_s <= ma_l:
                # 卖出
                if pre_price > 0 and entry_price > 0:
                    return_pct = (current_price / entry_price - 1) * position_ratio
                    equity.append(equity[-1] * (1 + (current_price / pre_price - 1) * position_ratio))
                    trade_log.append({
                        'datetime': row["datetime"],
                        'date': row["date"],
                        'action': 'SELL_ALL',
                        'price': current_price,
                        'return2': return_pct
                    })
                else:
                    equity.append(equity[-1])
                
                position_ratio = 0
                entry_price = 0
                
            elif position_ratio > 0 and pre_price > 0:
                # 持有仓位时更新净值
                equity.append(equity[-1] * (1 + (current_price / pre_price - 1) * position_ratio))
            else:
                equity.append(equity[-1])
            
            pre_price = current_price
        
        equity.pop(0)
        self.equity_curve_df = pd.DataFrame({
            "equity": equity,
            "datetime": truncate_btc_df["datetime"].reset_index(drop=True),
            "date": truncate_btc_df["date"].reset_index(drop=True)
        })
        
        self.trade_log_df = pd.DataFrame(trade_log)
        if not self.trade_log_df.empty:
            self.trade_log_df["date"] = self.trade_log_df["datetime"].dt.date
        
        print(f"策略执行完成，共{len(self.trade_log_df)}笔交易")
        
    def calculate_performance(self, nav_series, nav_benchmark, trade_log_df):
        """计算策略表现"""
        performance_obj = Performance()
        
        if nav_series.max() == nav_series.min() == 1 and trade_log_df.empty:
            return Performance()
        
        daily_returns = nav_series.pct_change().dropna()
        benchmark_returns = nav_benchmark.pct_change().dropna()
        
        total_return = nav_series.iloc[-1] / nav_series.iloc[0] - 1
        n_days = len(nav_series)
        
        annualized_return = (1 + total_return) ** (self.trading_days_one_year / n_days) - 1
        annualized_return = annualized_return * 100
        performance_obj.annualized_return = annualized_return
        
        performance_obj.cumulative_return = total_return * 100
        
        annualized_volatility = daily_returns.std() * np.sqrt(self.trading_days_one_year) * 100
        performance_obj.annualized_volatility = annualized_volatility
        
        annualized_sharpe = (annualized_return - self.risk_free_rate) / annualized_volatility if annualized_volatility != 0 else 0
        performance_obj.sharpe = annualized_sharpe
        performance_obj.annualized_sharpe = annualized_sharpe
        
        peak = nav_series.expanding().max()
        drawdown = (nav_series - peak) / peak
        max_drawdown = abs(drawdown.min()) * 100
        performance_obj.max_drawdown = max_drawdown
        
        benchmark_annualized_return = (nav_benchmark.iloc[-1] / nav_benchmark.iloc[0]) ** (
            self.trading_days_one_year / n_days) - 1
        benchmark_annualized_return = benchmark_annualized_return * 100
        excess_return = annualized_return - benchmark_annualized_return
        performance_obj.excess_return = excess_return
        
        excess_returns = daily_returns.reset_index(drop=True) - benchmark_returns.reset_index(drop=True)
        excess_cumret = (1 + excess_returns).cumprod()
        cumulative_excess = (excess_cumret.iloc[-1] - 1) * 100
        performance_obj.excess_cumulative = cumulative_excess
        
        calmar = annualized_return / abs(max_drawdown) if max_drawdown != 0 else 0
        performance_obj.calmar = calmar
        
        if not trade_log_df.empty:
            sell_count = len(trade_log_df[trade_log_df['action'].str.contains('SELL', na=False)])
            total_trades = sell_count
            yearly_trades = total_trades / (n_days / 365)
            performance_obj.total_trades = total_trades
            performance_obj.yearly_trades = yearly_trades
            performance_obj.yearly_count = yearly_trades
            
            sell_df = trade_log_df[trade_log_df['action'].str.contains('SELL', na=False)]
            if 'return2' in sell_df.columns and not sell_df.empty:
                winning_count = len(sell_df[sell_df['return2'] > 0])
                performance_obj.winning_rate = (winning_count / len(sell_df)) * 100 if len(sell_df) > 0 else 0
            else:
                performance_obj.winning_rate = 0
        else:
            performance_obj.total_trades = 0
            performance_obj.yearly_trades = 0
            performance_obj.winning_rate = 0
        
        return performance_obj
        
    def run_backtest(self):
        """运行完整回测"""
        print("=" * 60)
        print("BTC交易策略回测开始")
        print("=" * 60)
        
        self.load_data()
        self.calculate_ma_alignment()
        self.execute_strategy()
        
        hold_btc_df = self.btc_df[self.btc_df['date'] >= self.backtest_begin_date]
        nav_strategy = self.equity_curve_df['equity'] / self.equity_curve_df['equity'].iloc[0]
        nav_benchmark = hold_btc_df['close'] / hold_btc_df['close'].iloc[0]
        
        self.performance_overall = self.calculate_performance(
            nav_strategy, nav_benchmark, self.trade_log_df
        )
        
        for year in range(2020, 2026):
            year_start = datetime.date(year, 1, 1)
            year_end = datetime.date(year, 12, 31)
            if year == 2025:
                year_end = datetime.date(2025, 9, 30)
            
            year_equity = self.equity_curve_df[
                (self.equity_curve_df['date'] >= year_start) & 
                (self.equity_curve_df['date'] <= year_end)
            ]
            
            year_btc = hold_btc_df[
                (hold_btc_df['date'] >= year_start) & 
                (hold_btc_df['date'] <= year_end)
            ]
            
            if len(year_equity) == 0:
                continue
                
            year_trade_log = self.trade_log_df[
                (self.trade_log_df['date'] >= year_start) & 
                (self.trade_log_df['date'] <= year_end)
            ]
            
            nav_strategy_year = year_equity['equity'] / year_equity['equity'].iloc[0]
            nav_benchmark_year = year_btc['close'] / year_btc['close'].iloc[0]
            
            self.performance_by_year[year] = self.calculate_performance(
                nav_strategy_year, nav_benchmark_year, year_trade_log
            )
        
        print("回测完成")
        
    def generate_html_report(self):
        """生成HTML报告"""
        print("生成HTML报告...")
        
        html = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>BTC交易策略回测报告</title>
    <style>
        body {{ font-family: Arial, sans-serif; margin: 20px; }}
        h1, h2 {{ color: #333; }}
        table {{ border-collapse: collapse; width: 100%; margin: 20px 0; }}
        th, td {{ border: 1px solid #ddd; padding: 8px; text-align: right; }}
        th {{ background-color: #4CAF50; color: white; }}
        tr:nth-child(even) {{ background-color: #f2f2f2; }}
    </style>
</head>
<body>
    <h1>BTC/USDT 交易策略回测报告</h1>
    
    <h2>策略参数</h2>
    <ul>
        <li>交易对：BTC/USDT</li>
        <li>交易周期：日线</li>
        <li>短均线：{self.short_ma}天SMA</li>
        <li>长均线：{self.long_ma}天SMA</li>
        <li>均线多头排列区间：{self.ma_range_start}-{self.ma_range_end}</li>
        <li>阈值：{self.bullish_threshold}</li>
        <li>回测区间：{self.backtest_begin_date} 至 {self.backtest_end_date}</li>
    </ul>
    
    <h2>全区间回测结果</h2>
    <table>
        <tr><th>指标</th><th>数值</th></tr>
        <tr><td>年化收益率</td><td>{self.performance_overall.annualized_return:.2f}%</td></tr>
        <tr><td>累计收益率</td><td>{self.performance_overall.cumulative_return:.2f}%</td></tr>
        <tr><td>超额收益率</td><td>{self.performance_overall.excess_return:.2f}%</td></tr>
        <tr><td>累计超额</td><td>{self.performance_overall.excess_cumulative:.2f}%</td></tr>
        <tr><td>最大回撤</td><td>{self.performance_overall.max_drawdown:.2f}%</td></tr>
        <tr><td>总交易次数</td><td>{self.performance_overall.total_trades:.0f}</td></tr>
        <tr><td>年化交易次数</td><td>{self.performance_overall.yearly_trades:.2f}</td></tr>
        <tr><td>Sharpe比率</td><td>{self.performance_overall.sharpe:.2f}</td></tr>
        <tr><td>Calmar比率</td><td>{self.performance_overall.calmar:.2f}</td></tr>
        <tr><td>胜率</td><td>{self.performance_overall.winning_rate:.2f}%</td></tr>
    </table>
    
    <h2>分年度回测结果</h2>
    <table>
        <tr>
            <th>年度</th><th>年化收益率</th><th>累计收益率</th><th>超额收益率</th>
            <th>最大回撤</th><th>Sharpe</th><th>Calmar</th><th>胜率</th>
        </tr>
"""
        
        for year, perf in sorted(self.performance_by_year.items()):
            html += f"""
        <tr>
            <td>{year}年</td>
            <td>{perf.annualized_return:.2f}%</td>
            <td>{perf.cumulative_return:.2f}%</td>
            <td>{perf.excess_return:.2f}%</td>
            <td>{perf.max_drawdown:.2f}%</td>
            <td>{perf.sharpe:.2f}</td>
            <td>{perf.calmar:.2f}</td>
            <td>{perf.winning_rate:.2f}%</td>
        </tr>
"""
        
        html += """    </table>
</body>
</html>"""
        
        with open('backtest_report.html', 'w', encoding='utf-8') as f:
            f.write(html)
        print("HTML报告已生成：backtest_report.html")
        
    def generate_excel_report(self):
        """生成Excel报告"""
        print("生成Excel报告...")
        
        with pd.ExcelWriter('backtest_report.xlsx', engine='openpyxl') as writer:
            overall_df = pd.DataFrame([self.performance_overall.dict()])
            overall_df.to_excel(writer, sheet_name='全区间结果', index=False)
            
            annual_data = []
            for year, perf in sorted(self.performance_by_year.items()):
                row = {'年度': year}
                row.update(perf.dict())
                annual_data.append(row)
            annual_df = pd.DataFrame(annual_data)
            annual_df.to_excel(writer, sheet_name='分年度结果')
            
            self.trade_log_df.to_excel(writer, sheet_name='交易日志')
            
        print("Excel报告已生成：backtest_report.xlsx")
        
    def plot_results(self):
        """绘制收益曲线"""
        print("绘制收益曲线...")
        
        fig, axes = plt.subplots(3, 1, figsize=(14, 10))
        
        hold_btc_df = self.btc_df[self.btc_df['date'] >= self.backtest_begin_date]
        nav_strategy = self.equity_curve_df['equity'] / self.equity_curve_df['equity'].iloc[0]
        nav_benchmark = hold_btc_df['close'] / hold_btc_df['close'].iloc[0]
        
        axes[0].plot(self.equity_curve_df['date'], nav_strategy * 100, label='Strategy', linewidth=2)
        axes[0].plot(hold_btc_df['date'], nav_benchmark * 100, label='Benchmark', linewidth=2)
        axes[0].set_title('Net Asset Value', fontsize=14)
        axes[0].set_ylabel('NAV')
        axes[0].legend()
        axes[0].grid(True, alpha=0.3)
        
        peak = nav_strategy.expanding().max()
        drawdown = (nav_strategy - peak) / peak * 100
        axes[1].fill_between(self.equity_curve_df['date'], 0, drawdown, alpha=0.3, color='red')
        axes[1].plot(self.equity_curve_df['date'], drawdown, label='Drawdown', linewidth=1, color='darkred')
        axes[1].set_title('Drawdown', fontsize=14)
        axes[1].set_ylabel('Drawdown (%)')
        axes[1].legend()
        axes[1].grid(True, alpha=0.3)
        
        excess_returns = (nav_strategy.reset_index(drop=True) - nav_benchmark.reset_index(drop=True)) * 100
        axes[2].plot(self.equity_curve_df['date'], excess_returns, label='Excess Return', linewidth=2, color='green')
        axes[2].set_title('Excess Return', fontsize=14)
        axes[2].set_ylabel('Excess Return (%)')
        axes[2].set_xlabel('Date')
        axes[2].legend()
        axes[2].grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('backtest_chart.png', dpi=300, bbox_inches='tight')
        print("图表已保存：backtest_chart.png")
        plt.close()
        
    def save_trade_log(self):
        """保存交易日志"""
        self.trade_log_df.to_csv('trade_log.csv', index=False)
        print("交易日志已保存：trade_log.csv")


def main():
    strategy = BTCTradingStrategy()
    strategy.run_backtest()
    strategy.generate_html_report()
    strategy.generate_excel_report()
    strategy.plot_results()
    strategy.save_trade_log()
    
    print("\n" + "=" * 60)
    print("所有报告已生成完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()

